# Copyright (c) Microsoft. All rights reserved.
# Licensed under the MIT license. See LICENSE file in the project root for full license information.

"""
Version 3.0
what's new:
- support 3 layers
"""

import numpy as np
import time
import math
import os
import sys

from HelperClass2.HyperParameters_3_0 import *
from HelperClass2.MnistImageDataReader import *
from HelperClass2.TrainingHistory_2_3 import *
from HelperClass2.LossFunction_1_1 import *
from HelperClass2.ActivatorFunction_2_0 import *
from HelperClass2.ClassifierFunction_2_0 import *
from HelperClass2.WeightsBias_1_0 import *

class NeuralNet_3_0(object):
    def __init__(self, hp, model_name):
        """
        初始化神经网路，三层神经网络
        :param hp: 包含该神经网络模型所需各类参数，包括输入层神经元个数、隐藏层神经元个数、学习率、批大小等等
        :param model_name: 模型文件名
        """
        self.hp = hp
        self.model_name = model_name
        self.subfolder = os.getcwd() + "/" + self.__create_subfolder()
        print(self.subfolder)
        # 初始化输入层到第一个隐藏层的权重和偏置
        self.wb1 = WeightsBias_1_0(self.hp.num_input, self.hp.num_hidden1, self.hp.init_method, self.hp.eta)
        self.wb1.InitializeWeights(self.subfolder, False)
        # 初始化第一个隐藏层到第二个隐藏层的权重和偏置
        self.wb2 = WeightsBias_1_0(self.hp.num_hidden1, self.hp.num_hidden2, self.hp.init_method, self.hp.eta)
        self.wb2.InitializeWeights(self.subfolder, False)
        # 初始化第二个隐藏层到输出层的权重和偏置
        self.wb3 = WeightsBias_1_0(self.hp.num_hidden2, self.hp.num_output, self.hp.init_method, self.hp.eta)
        self.wb3.InitializeWeights(self.subfolder, False)

    def __create_subfolder(self):
        if self.model_name != None:
            path = self.model_name.strip()
            path = path.rstrip("/")
            isExists = os.path.exists(path)
            if not isExists:
                os.makedirs(path)
            return path

    def forward(self, batch_x):
        """
        前向计算函数，batch_x作为输入层，依次经过隐藏层和输出层，计算得到该神经网络的输出
        :param batch_x: 输入层
        :return:
        """
        # batch_x作为输入层，计算第一个隐藏层的输出
        self.Z1 = np.dot(batch_x, self.wb1.W) + self.wb1.B
        # 计算第一个隐藏层的输出经过Sigmoid函数激活后的输出
        self.A1 = Sigmoid().forward(self.Z1)
        # 这行被注释的公式也就是上一行的Sigmoid前向计算的公式
        # self.A1 = 1.0 / (1.0 + np.exp(-self.Z1))
        # 第一个隐藏层激活后的输出作为第二个隐藏层的输入层，计算第二个隐藏层的输出
        self.Z2 = np.dot(self.A1, self.wb2.W) + self.wb2.B
        # 计算第二个隐藏层的输出经过Tanh激活函数激活后的输出
        self.A2 = Tanh().forward(self.Z2)
        # 这行被注释的公式也就是上一行Tanh激活函数前向计算的公式
        # self.A2 = 2.0 / (1.0 + np.exp(-2*self.Z2)) - 1.0
        # 第二个隐藏层激活后的输出作为输出层的输入层，计算输出层的输出结果
        self.Z3 = np.dot(self.A2, self.wb3.W) + self.wb3.B
        # 如果我们要的效果是二分类任务，就对输出层的结果应用Logistic激活函数，作为输出层激活后的输出
        if self.hp.net_type == NetType.BinaryClassifier:
            self.A3 = Logistic().forward(self.Z3)
        # 如果要的效果的多重分类任务，就多输出层的结果应用SoftMax激活函数，作为输出层激活后的输出
        elif self.hp.net_type == NetType.MultipleClassifier:
            self.A3 = Softmax().forward(self.Z3)
        # 如果是其他类型比如简单线性拟合任务，那这里就不需要激活，直接将输出层的结果作为该模型最终的输出
        else:   # NetType.Fitting
            self.A3 = self.Z3
        # 将输出层的最终结果作为模型的输出
        self.output = self.A3

    def backward(self, batch_x, batch_y):
        """
        反向传播函数，计算出各层权重和偏置的梯度，为后续的权重更新操作提供依据
        :param batch_x:
        :param batch_y:
        :return:
        """
        # 批量下降，需要除以样本数量，否则会造成梯度爆炸
        # 获取输入数据的样本数量
        m = batch_x.shape[0]

        # 计算输出层的误差项dz3
        dZ3 = self.output - batch_y
        # 根据链式法则，计算输出层权重的梯度
        self.wb3.dW = np.dot(self.A2.T, dZ3)/m
        # 计算输出层的偏置的梯度
        self.wb3.dB = np.sum(dZ3, axis=0, keepdims=True)/m

        # 计算第二个隐藏层的误差传播项dA2
        dA2 = np.dot(dZ3, self.wb3.W.T)
        # 因为这里在前向计算的时候经过了激活函数Tanh处理，所以这里要使用tanh的反向计算还原第二个隐藏层的误差项dZ2
        dZ2,_ = Tanh().backward(None, self.A2, dA2)
        #da = 1 - np.multiply(self.A2, self.A2)
        #dZ2 = np.multiply(dA2, da)
        # 根据链式法则计算第二个隐藏层的权重的梯度
        self.wb2.dW = np.dot(self.A1.T, dZ2)/m
        # 计算第二个隐藏层的偏置的梯度
        self.wb2.dB = np.sum(dZ2, axis=0, keepdims=True)/m

        # 计算第一个隐藏层的误差项dA1
        dA1 = np.dot(dZ2, self.wb2.W.T) 
        # 这里同理，利用Sigmoid的反向传播方法还原第一个隐藏层的误差项dZ1
        dZ1,_ = Sigmoid().backward(None, self.A1, dA1)
        #da = np.multiply(self.A1, 1-self.A1)
        #dZ1 = np.multiply(dA1, da)

        # 根据链式法则，计算第一个隐藏层权重的梯度
        self.wb1.dW = np.dot(batch_x.T, dZ1)/m
        # 计算第一个隐藏层偏置的梯度
        self.wb1.dB = np.sum(dZ1, axis=0, keepdims=True)/m

    def update(self):
        self.wb1.Update()
        self.wb2.Update()
        self.wb3.Update()

    def inference(self, x):
        self.forward(x)
        return self.output

    def train(self, dataReader, checkpoint, need_test):
        """
        神经网络的争个训练过程和打印训练花费时间，包括前向计算、反向传播、更新权重和偏置、检查、测试等
        :param dataReader: 数据读取对象
        :param checkpoint: 检查点
        :param need_test: 是否需要测试
        :return:
        """
        t0 = time.time()  # 记录开始训练的时间点
        # 创建一个用于记录训练过程损失值和准确率变化历史的对象
        self.loss_trace = TrainingHistory_2_3()
        # 根据网络类型创建相应的损失函数对象
        self.loss_func = LossFunction_1_1(self.hp.net_type)
        loss = 10
        if self.hp.batch_size == -1:
            self.hp.batch_size = dataReader.num_train
        # 计算最大迭代次数，这里只每一轮训练的最大迭代次数，最大训练轮次为max_epoch
        max_iteration = math.ceil(dataReader.num_train / self.hp.batch_size)
        # 根据检查点比例计算出每多少次迭代后进行一次检查（计算损失和准确率）
        checkpoint_iteration = int(max_iteration * checkpoint)
        need_stop = False  # 初始化一个标记，表示是否需要停止训练（该模型是否已经达标）
        # 开始每个轮次的训练
        for epoch in range(self.hp.max_epoch):
            # 对数据集进行随机打乱操作
            dataReader.Shuffle()
            # 开始每个轮次的每次迭代训练，一次训练batch_size个数据
            for iteration in range(max_iteration):
                # 根据数据读取对象获取该批次的输入数据batch_x和对应真实标签batch_y
                batch_x, batch_y = dataReader.GetBatchTrainSamples(self.hp.batch_size, iteration)
                # batch_x经过神经网络的前向计算得出神经网络的输出
                self.forward(batch_x)
                # 执行神经网络的反向传播函数，计算每一层的权重和偏置的梯度（更新量）
                self.backward(batch_x, batch_y)
                # 更新每一层的权重和偏置
                self.update()
                # 计算当前是总的多少次迭代，用于看是否要检查
                total_iteration = epoch * max_iteration + iteration
                # 如果当前总的迭代次数加1（因为从0开始）后能别检查点迭代次数整除（也就是达到了检查点），那就需要在该次迭代中执行检查
                if (total_iteration+1) % checkpoint_iteration == 0:
                    # 执行检查函数，该函数记录在当前迭代时模型的损失和准确率，最终返回模型损失是否达标
                    need_stop = self.CheckErrorAndLoss(dataReader, batch_x, batch_y, epoch, total_iteration)
                    if need_stop:
                        break
            if need_stop:
                break
        # 保存训练结果（各层的权重和偏置）
        self.SaveResult()
        # 记录训练结束的时间点
        t1 = time.time()
        # 输出本次训练总共花费时间
        print("time used:", t1 - t0)
        # 如果需要测试，就执行测试函数并打印测试模型的正确率
        if need_test:
            print("testing...")
            accuracy = self.Test(dataReader)
            print(accuracy)

    def CheckErrorAndLoss(self, dataReader, train_x, train_y, epoch, total_iteration):
        """
        检查训练过程中的损失和准确率情况，包括训练集和验证集，并判断模型是否达标
        :param dataReader: 数据读取对象
        :param train_x: 批次输入数据
        :param train_y: 批次输入数据的对应的真实标签
        :param epoch: 当前训练的轮次
        :param total_iteration: 当前总的迭代次数
        :return: 是否需要停止训练（模型是否达标）
        """
        print("epoch=%d, total_iteration=%d" % (epoch, total_iteration))

        # 对于当前批次的输入对应的模型的输出
        self.forward(train_x)
        # 计算训练集的损失
        loss_train = self.loss_func.CheckLoss(self.output, train_y)
        # 计算训练集的准确率
        accuracy_train = self.__CalAccuracy(self.output, train_y)
        print("loss_train=%.6f, accuracy_train=%f" %(loss_train, accuracy_train))

        # 获取验证集
        vld_x, vld_y = dataReader.GetValidationSet()
        # 计算验证集中输入通过模型后的验证集对应预测值
        self.forward(vld_x)
        # 计算验证集损失
        loss_vld = self.loss_func.CheckLoss(self.output, vld_y)
        # 计算验证集的准确率
        accuracy_vld = self.__CalAccuracy(self.output, vld_y)
        print("loss_valid=%.6f, accuracy_valid=%f" %(loss_vld, accuracy_vld))
        # 将当前轮次、总迭代次数、训练集损失和准确率、验证集损失和准确率等信息添加到记录中
        need_stop = self.loss_trace.Add(epoch, total_iteration, loss_train, accuracy_train, loss_vld, accuracy_vld)
        # 判断验证集的损失是否小于等于我们预设定的值eps，如果满足则可以停止训练
        if loss_vld <= self.hp.eps:
            need_stop = True
        return need_stop

    def Test(self, dataReader):
        """
        从数据集中抽取一部分作为测试集，测试该模型的准确率
        :param dataReader: 数据读取对象
        :return:
        """
        x, y = dataReader.GetTestSet()
        self.forward(x)
        correct = self.__CalAccuracy(self.output, y)
        print(correct)

    def __CalAccuracy(self, a, y):
        """
        根据网络类型计算准确率
        :param a: 神经网络预测值
        :param y: 对应的真实标签值
        :return: 该神经网络的准确率
        """
        # 断言确保预测值和真实标签的个数相同
        assert(a.shape == y.shape)
        m = a.shape[0]
        # 如果时简单拟合任务，就用方差和均方误差计算决定系数r2，r2越接近1表示拟合效果越好
        if self.hp.net_type == NetType.Fitting:
            var = np.var(y)
            mse = np.sum((a-y)**2)/m
            r2 = 1 - mse / var
            return r2
        # 如果是二分类任务，就计算正确预测类别的个数也就是准确率
        elif self.hp.net_type == NetType.BinaryClassifier:
            b = np.round(a)
            r = (b == y)
            correct = np.sum(r)
            return correct/m
        # 如果是多重分类任务，也是计算准确率
        elif self.hp.net_type == NetType.MultipleClassifier:
            ra = np.argmax(a, axis=1)
            ry = np.argmax(y, axis=1)
            r = (ra == ry)
            correct = np.sum(r)
            return correct/m

    def SaveResult(self):
        """
        把最终的训练结果（每层的权重的偏置）保存到文件
        :return:
        """
        self.wb1.SaveResultValue(self.subfolder, "wb1")
        self.wb2.SaveResultValue(self.subfolder, "wb2")
        self.wb3.SaveResultValue(self.subfolder, "wb3")

    def LoadResult(self):
        """
        加载模型
        :return:
        """
        self.wb1.LoadResultValue(self.subfolder, "wb1")
        self.wb2.LoadResultValue(self.subfolder, "wb2")
        self.wb3.LoadResultValue(self.subfolder, "wb3")

    def ShowTrainingHistory(self, xcoord):
        """
        可视化训练过程中损失和准确率变化
        :return:
        """
        self.loss_trace.ShowLossHistory(self.hp, xcoord)

    def GetTrainingTrace(self):
        """
        获取训练的损失和准确率历史记录实例对象
        :return: 训练的损失和准确率历史记录实例对象
        """
        return self.loss_trace

    def GetEpochNumber(self):
        """
        获取当前训练的轮次
        :return: 当前训练的轮次
        """
        return self.loss_trace.GetEpochNumber()

    def GetLatestAverageLoss(self, count=10):
        """
        获取最后count次的验证集平均损失（count>保存的个数则返回最后一次损失值）
        :param count: 获取最后count次的验证集平均损失
        :return: 最后count次的验证集平均损失（count>保存的个数则返回最后一次损失值）
        """
        return self.loss_trace.GetLatestAverageLoss(count)

    def DumpLossHistory(self, filename):
        return self.loss_trace.Dump(filename)
